package unreleased.geppetto.cat.models;


import geppetto.cat.alignments.Alignment;
import geppetto.cat.alignments.AlignmentEvaluator;
import geppetto.cat.alignments.AlignmentsSet;
import geppetto.cat.alignments.AlignmentEvaluator.Evaluation;
import geppetto.cat.common.StaticTools;
import geppetto.cat.constrains.ConstrainedProjectionStats;
import geppetto.cat.corpus.BilingualCorpus;
import geppetto.cat.models.AbstractModel;
import geppetto.cat.models.AbstractSparseTranslationTable;
import geppetto.cat.models.AgreementM1;
import geppetto.cat.models.HMM;
import geppetto.cat.models.stats.AgreeEStepStats;
import geppetto.cat.models.stats.ProjectionStats;

import java.io.BufferedReader;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.FileReader;
import java.io.IOException;
import java.io.PrintStream;




public class AgreementHMMWithClone extends AbstractModel {

	HMM forward;
	HMM backward;
	double epsilon;
	int _projectionIterations;
	boolean _trained = false;
	

	public AgreementHMMWithClone(BilingualCorpus forwardCorpus,
			BilingualCorpus backwardCorpus, double epsilon, int projectionIterations) {
		this(forwardCorpus, backwardCorpus, null, null, epsilon, projectionIterations);
	}
	
	public AgreementHMMWithClone(BilingualCorpus forwardCorpus,
			BilingualCorpus backwardCorpus, double epsilon, double smooth, int projectionIterations) {
		this(forwardCorpus, backwardCorpus, null, null, epsilon,smooth, projectionIterations);
	}

	public AgreementHMMWithClone(BilingualCorpus forwardCorpus,
			BilingualCorpus backwardCorpus, AbstractSparseTranslationTable forwardT,
			AbstractSparseTranslationTable backwardT, double epsilon, int projectionIterations) {
			this(forwardCorpus, backwardCorpus, forwardT, backwardT, epsilon, 0, projectionIterations);
	}
	
	public AgreementHMMWithClone(BilingualCorpus forwardCorpus,
			BilingualCorpus backwardCorpus, AbstractSparseTranslationTable forwardT,
			AbstractSparseTranslationTable backwardT, double epsilon, double smooth,int projectionIterations) {
		this.epsilon = epsilon;
		this._projectionIterations = projectionIterations;
		this._corpus = forwardCorpus;
		System.out.println("Init Agreement HMM with epsilon " + epsilon + " smooth " + smooth + " projectionIterations " + projectionIterations );
		if(forwardT == null){
			forward = new HMM(forwardCorpus,smooth);
		}else{
			forward = new HMM(forwardCorpus,forwardT,smooth);
		}
		if(backwardT == null){
			backward = new HMM(backwardCorpus,smooth);
		}else{
			backward = new HMM(backwardCorpus, backwardT,smooth);
	
		}
	}
	AgreementHMMWithClone() {

	}

	// /////////////////SAVE AND LOAD MODELS

	public void saveModel(String directory) {
		forward.saveModel(directory + "forward");
		backward.saveModel(directory + "backward");
		try {
			PrintStream file = new PrintStream(new FileOutputStream(directory+ "epsilon"));
			file.println(_projectionIterations);
			file.println(epsilon);
		} catch (FileNotFoundException e) {
			System.out.println(getName()+": Could not save model");
			System.exit(-1);
		}
	}

	public static AgreementHMMWithClone loadModel(BilingualCorpus forwardCorpus,
			BilingualCorpus backwardCorpus, String directory) {
		AgreementHMMWithClone model = new AgreementHMMWithClone();
		try {
			BufferedReader file = new BufferedReader(new FileReader(directory
					+ "epsilon"));
			model._projectionIterations = Integer.parseInt(file.readLine());
			model.epsilon = Double.parseDouble(file.readLine());
			System.out.println("Reading Forward: " + directory + "forward");
			model.forward = HMM.loadModel(forwardCorpus, directory + "forward");
			model._corpus = forwardCorpus;
			System.out.println("Reading Forward: " + directory + "backward");
			model.backward = HMM.loadModel(backwardCorpus, directory
					+ "backward");
		} catch (FileNotFoundException e) {
			System.out
					.println(model.forward.getName() + " Could not read model file not found "
							+ e.toString());
			System.exit(-1);
		} catch (NumberFormatException e) {
			System.out
					.println(model.forward.getName() + "Could not read model number conversion exception"
							+ e.toString());
			System.exit(-1);

		} catch (IOException e) {
			System.out
					.println(model.forward.getName() + ": Could not read model"
							+ e.toString());
			System.exit(-1);
		}
		return model;
	}

	AgreementHMMWithClone(HMM forward, HMM backward) {
		this.forward = forward;
		this.backward = backward;
		this._corpus = forward._corpus;
	}

	public AgreementHMMWithClone getBackwardModel() {
		AgreementHMMWithClone res = new AgreementHMMWithClone(backward, forward);
		res._trained = _trained;
		return res;
	}

	public String getName() {
		return "Agremement HMM";
	}
	
	public void setTrained() {
		_trained = true;
		forward.setTrained();
		backward.setTrained();
	}

	public void initializeTrain() {
		forward.initializeTrain();
		backward.initializeTrain();
	}

	public void finalizeTrain() {
		forward.finalizeTrain();
		backward.finalizeTrain();
	}

	
	public double getNullPhrasePosterior(int phraseNumber, byte phraseSource,
			int[] foreingSentence, int[] sourceSentence, int startForeignIndex,
			int endForeignIndex) {
		return forward.getNullPhrasePosterior(phraseNumber, phraseSource,
				foreingSentence, sourceSentence, startForeignIndex,
				endForeignIndex);
	}

	public double getPhrasePosterior(int phraseNumber, byte phraseSource,
			int[] foreingSentence, int[] sourceSentence, int startSourceIndex,
			int endSourceIndex, int startForeignIndex, int endForeignIndex) {
		return forward.getPhrasePosterior(phraseNumber, phraseSource,
				foreingSentence, sourceSentence, startSourceIndex,
				endSourceIndex, startForeignIndex, endForeignIndex);
	}

	public ProjectionStats projectPosteriors(double[][][] probCache,
			double[][][] forwardTables, double[][][] backwardTables,
			double [][][] posteriors, int sSize, int fSize, int sentenceNumber
			/*, boolean postStats*/) {

		ProjectionStats stats = new ProjectionStats(_projectionIterations);
		double posteriorsF[][] = posteriors[0];double posteriorsB[][] = posteriors[1];
		double[][] originalPosteriorsF = StaticTools.copyMatrix(posteriorsF,  fSize,sSize*2);
		double[][] originalPosteriorsB = StaticTools.copyMatrix(posteriorsB,  sSize,fSize*2);
		double probCacheF[][] = probCache[0];double probCacheB[][] = probCache[1];
		double forwardF[][] = forwardTables[0];double forwardB[][] = forwardTables[1];
		double backwardF[][] = backwardTables[0];double backwardB[][] = backwardTables[1];
		
		double[][] probCacheFP = probCacheF.clone();
		double[][] probCacheBP = probCacheB.clone();
		
		//double[][] probCacheFP = StaticTools.copyMatrix(probCacheF,  fSize,sSize*2);
		//double[][] probCacheBP = StaticTools.copyMatrix(probCacheB,  sSize,fSize*2);
		
		double[][] phi = new double[fSize][sSize];
		for (int grI = 0; grI < _projectionIterations; grI++) {
			double lRate = 10.0 / (10.0 + grI);
			double totalViolation = 0;
			for (int fi = 0; fi < fSize; fi++) {
				for (int si = 0; si < sSize; si++) {
					// Update phi and lambda
					double ePsi = posteriorsF[fi][si] - posteriorsB[si][fi];
					double phiUpdate = lRate * (-ePsi);
					totalViolation += Math.abs(-ePsi);
					double newlambda = Math.abs(phi[fi][si]) - lRate * epsilon;
					double newPhi = phiUpdate + phi[fi][si];
					// Project
					if (-newPhi <= newlambda && newlambda <= newPhi) {
						newPhi = (newPhi + newlambda) / 2;
						newlambda = newPhi;
					} else if (newPhi <= newlambda && newlambda <= -newPhi) {
						newPhi = (newPhi - newlambda) / 2;
						newlambda = -newPhi;
					} else {
						newPhi = 0;
						newlambda = 0;
					}
					probCacheF[fi][si] = probCacheFP[fi][si] * Math.exp(newPhi);
					probCacheB[si][fi] = probCacheBP[si][fi] * Math.exp(-newPhi);
					phi[fi][si] = newPhi;
				}
			}
			// Renormalizing
			forwardF = forward.makeForward(sSize, fSize, probCacheF);
			backwardF = forward.makeBackward(sSize, fSize, probCacheF);
			double likelihoodF = forward.makeLikelihood(forwardF[0],
					backwardF[0]);
			posteriorsF = forward.makePosterior(forwardF, backwardF,
					likelihoodF);

			
			forwardB = backward.makeForward(fSize, sSize, probCacheB);
			backwardB = backward.makeBackward(fSize, sSize, probCacheB);
			double likelihoodB = backward.makeLikelihood(forwardB[0],
					backwardB[0]);
			posteriorsB = backward.makePosterior(forwardB, backwardB,
					likelihoodB);

			
			//if (postStats) {			
			stats.addViolation(totalViolation, grI);
			//System.out.println("   "+sentenceNumber+"  final violation = "+totalViolation);	

				// Calculate thw KL(q|p) which is what we are minimizing
			stats.addForwardKL(geppetto.cat.common.StaticTools.KLDistance(posteriorsF, originalPosteriorsF, fSize,sSize * 2), grI);
			stats.addBackwardKL(geppetto.cat.common.StaticTools.KLDistance(posteriorsB, originalPosteriorsB, sSize,fSize * 2), grI);
			//}
			
		}
		posteriors[0] = posteriorsF;posteriors[1] = posteriorsB;
		forwardTables[0] = forwardF; forwardTables[1] = forwardB;
		backwardTables[0] = backwardF; backwardTables[1] = backwardB;
		probCache[0] = probCacheF; probCache[1] = probCacheB;
		return stats;
	}
		
	
//	 Its here just because of the stats. Should be removed.
	public void train(int iterations) {
		_numberTrainingIterations = iterations;
		System.out.println("Starting " + getName() + " Training");
		initializeTrain();
		AgreeEStepStats d = new AgreeEStepStats(), old = new AgreeEStepStats();
		for (int k = 0; k < iterations; k++) {
			// /DEBUG CODE
			System.out.println("Iteration " + (k + 1));
			System.out.flush();
			// E-Step
			System.out.println(" e Step");
			d.startTime();
			d.add(eStep());
			
			d.stopTime();
			System.out.println(d.makeVerbose(old));
			// M-Step
			old = d;
			d = new AgreeEStepStats();
			mStep();
			System.out.println("m Step");
			System.out.flush();
		}
		setTrained();
		finalizeTrain();
		System.out.println();
	}
	
	public AgreeEStepStats eStep(/*boolean postStats*/) {
		double totalLikelihoodF = 0;
		double totalLikelihoodB = 0;
		ProjectionStats totalPStats = new ProjectionStats(_projectionIterations);
		forward.clearCounts();
		backward.clearCounts();
		for (int i = 0; i < forward._nSentences; i++) {
			int[] s = forward._corpus.getSourceSentence(i,
					BilingualCorpus.TRAIN_CORPUS);
			final int sSize = s.length;
			int[] f = forward._corpus.getForeignSentence(i,
					BilingualCorpus.TRAIN_CORPUS);
			final int fSize = f.length;


			// indexed by[fi][si] sSize..2*sSize corresponds to null word
			// probability
			double[][] probCacheF = forward.makeProbCache(f, s);
			double[][] forwardF = forward.makeForward(sSize, fSize, probCacheF);
			double[][] backwardF = forward.makeBackward(sSize, fSize,
					probCacheF);
			double likelihoodF = forward.makeLikelihood(forwardF[0],
					backwardF[0]);

			// indexed by[si][fi] sSize..2*sSize corresponds to null word
			// probability
			double[][] probCacheB = backward.makeProbCache(s, f);
			double[][] forwardB = backward
					.makeForward(fSize, sSize, probCacheB);
			double[][] backwardB = backward.makeBackward(fSize, sSize,
					probCacheB);
			double likelihoodB = backward.makeLikelihood(forwardB[0],
					backwardB[0]);

			assert forwardB.length == sSize;
			assert backwardB.length == sSize;
			double[][] posteriorsF = forward.makePosterior(forwardF, backwardF,
					likelihoodF);
			double[][] posteriorsB = backward.makePosterior(forwardB,
					backwardB, likelihoodB);
			
			
//			 check for underflow
			if (likelihoodF <= 1.0e-200 || likelihoodB <= 1.0e-200){ 
				_numericUnderFlow++;
				continue;
			}
			double [][][] probCache = {probCacheF,probCacheB};
			double [][][] forwardTables = {forwardF,forwardB};
			double [][][] backwardTables = {backwardF,backwardB};
			double [][][] posteriors = {posteriorsF,posteriorsB};
	
			ProjectionStats stats = projectPosteriors(probCache,forwardTables,backwardTables,posteriors,sSize, fSize,i); 
			totalPStats.add(stats);
			
			probCacheF = probCache[0];probCacheB = probCache[1];
			forwardF = forwardTables[0];forwardB = forwardTables[1];
			backwardF = backwardTables[0];backwardB = backwardTables[1];
			posteriorsF = posteriors[0];posteriorsB = posteriors[1];
			
			likelihoodF = forward.makeLikelihood(forwardF[0], backwardF[0]);
			likelihoodB = backward.makeLikelihood(forwardB[0], backwardB[0]);
			totalLikelihoodF += Math.log(likelihoodF);
			totalLikelihoodB += Math.log(likelihoodB);
			
			forward.addToCounts(s, f, posteriorsF);
			backward.addToCounts(f, s, posteriorsB);

			forward.addToTransitions(probCacheF, forwardF, backwardF,
					likelihoodF);
			backward.addToTransitions(probCacheB, forwardB, backwardB,
					likelihoodB);
		}
		AgreeEStepStats d = new AgreeEStepStats();
		d.logLikelihoodF = totalLikelihoodF;
		d.logLikelihoodB = totalLikelihoodB;
		d.numSents = forward._corpus.getNumberOfTrainingSentences();
		d.pStats = totalPStats;
		return d;
	}

	public void mStep() {
		forward.mStep();
		backward.mStep();
	}

	public Alignment posteriorDecodingAlignment(int sentenceNumber,
			byte sentenceSource, float treshhold,boolean projectPosteriors, ConstrainedProjectionStats stats) {
		return forward.posteriorDecodingAlignment(sentenceNumber,
				sentenceSource, treshhold);
	}

	public Alignment viterbiAlignment(int sentenceNumber, byte sentenceSource, boolean projectPosteriors, ConstrainedProjectionStats stats) {
		if(projectPosteriors){
			System.out.println("Ignoring project posteriors");
		}
		return forward.viterbiAlignment(sentenceNumber, sentenceSource, projectPosteriors);
	}
	
	/*
	public void projectionStep(double[][] probCacheF, double[][] probCacheB,
			int sSize, int fSize) {

		double[][] forwardF = forward.makeForward(sSize, fSize, probCacheF);
		double[][] backwardF = forward.makeBackward(sSize, fSize, probCacheF);
		double likelihoodF = forward.makeLikelihood(forwardF[0], backwardF[0]);

		// indexed by[si][fi] sSize..2*sSize corresponds to null word
		// probability
		double[][] forwardB = backward.makeForward(fSize, sSize, probCacheB);
		double[][] backwardB = backward.makeBackward(fSize, sSize, probCacheB);
		double likelihoodB = backward.makeLikelihood(forwardB[0], backwardB[0]);

		// Original caches
		double[][] probCacheFP = cat.common.StaticTools.copyMatrix(probCacheF,
				fSize, sSize * 2);
		double[][] probCacheBP = cat.common.StaticTools.copyMatrix(probCacheB,
				sSize, fSize * 2);

		double[][] posteriorsF = forward.makePosterior(forwardF, backwardF,
				likelihoodF);
		double[][] posteriorsB = backward.makePosterior(forwardB, backwardB,
				likelihoodB);

		double[][] phi = new double[fSize][sSize];
		for (int grI = 0; grI < _projectionIterations; grI++) {

			double lRate = 10.0 / (10.0 + grI);
			double totalViolation = 0;
			for (int fi = 0; fi < fSize; fi++) {
				for (int si = 0; si < sSize; si++) {
					// Update phi and lambda
					double ePsi = posteriorsF[fi][si] - posteriorsB[si][fi];
					double phiUpdate = lRate * (-ePsi);
					totalViolation += Math.abs(-ePsi);
					double newlambda = Math.abs(phi[fi][si]) - lRate * epsilon;
					double newPhi = phiUpdate + phi[fi][si];
					// Project
					if (-newPhi <= newlambda && newlambda <= newPhi) {
						newPhi = (newPhi + newlambda) / 2;
						newlambda = newPhi;
					} else if (newPhi <= newlambda && newlambda <= -newPhi) {
						newPhi = (newPhi - newlambda) / 2;
						newlambda = -newPhi;
					} else {
						newPhi = 0;
						newlambda = 0;
					}
					probCacheF[fi][si] = probCacheFP[fi][si] * Math.exp(newPhi);
					probCacheB[si][fi] = probCacheBP[si][fi]
							* Math.exp(-newPhi);
					phi[fi][si] = newPhi;
				}
			}
			// Renormalizing
			forwardF = forward.makeForward(sSize, fSize, probCacheF);
			backwardF = forward.makeBackward(sSize, fSize, probCacheF);
			likelihoodF = forward.makeLikelihood(forwardF[0], backwardF[0]);
			posteriorsF = forward.makePosterior(forwardF, backwardF,
					likelihoodF);

			forwardB = backward.makeForward(fSize, sSize, probCacheB);
			backwardB = backward.makeBackward(fSize, sSize, probCacheB);
			likelihoodB = backward.makeLikelihood(forwardB[0], backwardB[0]);
			posteriorsB = backward.makePosterior(forwardB, backwardB,
					likelihoodB);
			
			if (likelihoodF <= 1.0e-300 || likelihoodB <= 1.0e-300) {
				System.out.println("Oh no! likelhood became too small! -- failing this instance");
			}
		}
	}
*/
	
	/*
	public AlignmentsSet viterbiAlignmentsWithProjection(byte sentenceSource) {
		assert _trained;
		AlignmentsSet set = new AlignmentsSet();
		// System.out.println(_corpus.getNumSentences(sentenceSource));
		for (int i = 0; i < _corpus.getNumSentences(sentenceSource); i++) {
			set.addAlignment(viterbiAlignmentWithProjection(i, sentenceSource));
		}
		return set;
	}

	public Alignment viterbiAlignmentWithProjection(int sentenceNumber,
			byte sentenceSource) {
		int[] s = forward._corpus.getSourceSentence(sentenceNumber,
				sentenceSource);
		final int sSize = s.length;
		int[] f = forward._corpus.getForeignSentence(sentenceNumber,
				sentenceSource);
		final int fSize = f.length;
		double[][] probCacheF = forward.makeProbCache(f, s);
		double[][] probCacheB = backward.makeProbCache(s, f);
		projectionStep(probCacheF, probCacheB, sSize, fSize);
		double[][] forwardF = forward.makeForward(sSize, fSize, probCacheF);
		double[][] backwardF = forward.makeBackward(sSize, fSize, probCacheF);
		double likelihoodF = forward.makeLikelihood(forwardF[0], backwardF[0]);
		double[][] posteriorF = forward.makePosterior(forwardF, backwardF,
				likelihoodF);

		// Do viterbi but use the probCache instead of using the normal
		// translation probabilities

		// s.length..2*s.length is null word 
		double[][] totalScores = new double[s.length * 2][f.length];
		int[][] backPointers = new int[s.length * 2][f.length];

		// Fill position zero of posteriors
		for (int i = 0; i < s.length * 2; i++) {
			totalScores[i][0] = probCacheF[0][i];
		}

		// Fill the rest of the positions
		for (int foreignPos = 1; foreignPos < f.length; foreignPos++) {
			// s.length..2*s.length is null word 
			for (int nextSourceState = 0; nextSourceState < s.length * 2; nextSourceState++) {
				double maxScore = 0;
				int maxPosition = 0;
				// s.length..2*s.length is null word 
				for (int currentSourceState = 0; currentSourceState < s.length * 2; currentSourceState++) {
					double currScore;
					if (nextSourceState < s.length) {
						currScore = totalScores[currentSourceState][foreignPos - 1]
								* forward._distortion.getDistProb(
										nextSourceState, currentSourceState
												% s.length, s.length);
					} else {
						// s.length..2*s.length is null word 
						currScore = totalScores[currentSourceState][foreignPos - 1]
								* forward._distortion
										.getNullDistProb(nextSourceState
												- s.length, currentSourceState
												% s.length, s.length);
					}

					if (maxScore < currScore) {
						maxScore = currScore;
						maxPosition = currentSourceState;
					}
				}
				if (nextSourceState < s.length) {
					totalScores[nextSourceState][foreignPos] = probCacheF[foreignPos][nextSourceState]
							* maxScore;
				} else {
					// s.length..2*s.length is null word 
					totalScores[nextSourceState][foreignPos] = probCacheF[foreignPos][0]
							* maxScore;
				}
				backPointers[nextSourceState][foreignPos] = maxPosition;
			}
		}

		double max = 0;
		int max_ind = 0;
		// s.length..2*s.length is null word 
		for (int ind = 0; ind < s.length * 2; ind++) {
			double curr = totalScores[ind][f.length - 1];
			if (max < curr) {
				max = curr;
				max_ind = ind;
			}
		}

		int[] alignments = new int[f.length];
		Alignment a = new Alignment(sentenceNumber, sentenceSource, s.length,
				f.length);
		alignments[f.length - 1] = max_ind;
		for (int j = f.length - 2; j >= 0; j--) {
			alignments[j] = backPointers[alignments[j + 1]][j + 1];
		}
		for (int j = 0; j < f.length; j++) {
			// s.length..2*s.length is null word 
			if (alignments[j] < s.length) {
				a.add(alignments[j], j);
				a.addPosteriro(alignments[j], j, posteriorF[j][alignments[j]]);
			}
		}

		return a;

	}

	public Alignment posteriorDecodingAgreementAlignment(int sentenceNumber,
			byte sentenceSource, float treshhold) {
		int[] s = forward._corpus.getSourceSentence(sentenceNumber,
				sentenceSource);
		final int sSize = s.length;
		int[] f = forward._corpus.getForeignSentence(sentenceNumber,
				sentenceSource);
		final int fSize = f.length;

		
		 // indexed by[fi][si] sSize..2*sSize corresponds to null word
		 // probability
		 
		double[][] probCacheF = forward.makeProbCache(f, s);
		double[][] forwardF = forward.makeForward(sSize, fSize, probCacheF);
		double[][] backwardF = forward.makeBackward(sSize, fSize, probCacheF);
		double likelihoodF = forward.makeLikelihood(forwardF[0], backwardF[0]);

		double[][] probCacheB = backward.makeProbCache(s, f);
		double[][] forwardB = backward.makeForward(fSize, sSize, probCacheB);
		double[][] backwardB = backward.makeBackward(fSize, sSize, probCacheB);
		double likelihoodB = backward.makeLikelihood(forwardB[0], backwardB[0]);

		assert forwardB.length == sSize;
		assert backwardB.length == sSize;

		assert likelihoodB > 0 && likelihoodF > 0;

		
		 // indexed by[fi][si] sSize..2*sSize corresponds to null word
		 // probability
		 
		double[][] posteriorsF = forward.makePosterior(forwardF, backwardF,
				likelihoodF);
		
		double[][] posteriorsB = backward.makePosterior(forwardB, backwardB,
				likelihoodB);

		for (int fi = 0; fi < fSize; fi++) {
			for (int si = 0; si < sSize; si++) {
				posteriorsF[fi][si] = posteriorsF[fi][si] * posteriorsB[si][fi];
				posteriorsB[si][fi] = posteriorsF[fi][si];

			}
		}

		Alignment a = new Alignment(sentenceNumber, sentenceSource, sSize,
				fSize);

		for (int fi = 0; fi < fSize; fi++) {
			for (int si = 0; si < sSize; si++) {
				if (posteriorsF[fi][si] > treshhold)
					a.add(si, fi);
			}
		}
		return a;
	}
	*/
	
/*
	public void train(int iterations) {
		train(iterations, false);
	}
	*/

	// Normal EM training
	// TODO CODE IS THE SAME AS M1AGREE can just copy
	/*
	public void train(int iterations, boolean postStats) {
		System.out.println("Starting " + getName() + " Training");
		initializeTrain();
		AgreeEStepStats d = new AgreeEStepStats(), old = new AgreeEStepStats();
		for (int k = 0; k < iterations; k++) {

			// /DEBUG CODE

			System.out.println("Iteration " + (k + 1));
			System.out.flush();
			// E-Step
			System.out.println(" e Step");
			d.startTime();
			d.add(eStep(postStats));
			d.stopTime();
			System.out.println(d.makeVerbose(old));
			// M-Step
			old = d;
			d = new AgreeEStepStats();
			mStep();
			System.out.println("m Step");
			System.out.flush();
		}
		setTrained();
		finalizeTrain();
		System.out.println();
	}

	public Evaluation[] trainWithResults(int iterations, byte sentenceSource) {
		return trainWithResults(iterations, sentenceSource, false);
	}

	public Evaluation[] trainWithResults(int iterations, byte sentenceSource,
			boolean postStats) {
		Evaluation[] evals = new Evaluation[iterations];
		System.out.println("Starting " + getName() + " Training");
		initializeTrain();
		AgreeEStepStats d = new AgreeEStepStats(), old = new AgreeEStepStats();
		for (int k = 0; k < iterations; k++) {

			// /DEBUG CODE

			System.out.println("Iteration " + (k + 1));
			System.out.flush();
			// E-Step
			System.out.println(" e Step");
			d.startTime();
			d.add(eStep(postStats));
			d.stopTime();
			System.out.println(d.makeVerbose(old));
			// M-Step
			old = d;
			d = new AgreeEStepStats();
			mStep();
			System.out.println("m Step");
			System.out.flush();
			AlignmentsSet predicted = viterbiAlignments(sentenceSource);
			evals[k] = AlignmentEvaluator.evaluate(predicted, forward._corpus
					.getAlignments(sentenceSource));
		}
		setTrained();
		finalizeTrain();
		System.out.println();
		return evals;
	}
	*/
	
	
	/*

	// Prints a specific alignment before and after projection
	// These method should be abstracted
	public void printPosteriorAlignment(Alignment fa, Alignment ba,
			int sentenceNumber, byte sentenceSource,
			BilingualCorpus forwardCorpus, BilingualCorpus backwardCorpus)
			throws UnsupportedEncodingException {

		int[] s = forward._corpus.getSourceSentence(sentenceNumber,
				sentenceSource);
		final int sSize = s.length;
		int[] f = forward._corpus.getForeignSentence(sentenceNumber,
				sentenceSource);
		final int fSize = f.length;

		// indexed by[fi][si] sSize..2*sSize corresponds to null word
		// probability
		double[][] probCacheF = forward.makeProbCache(f, s);
		double[][] forwardF = forward.makeForward(sSize, fSize, probCacheF);
		double[][] backwardF = forward.makeBackward(sSize, fSize, probCacheF);
		double likelihoodF = forward.makeLikelihood(forwardF[0], backwardF[0]);

		// indexed by[si][fi] sSize..2*sSize corresponds to null word
		// probability
		double[][] probCacheB = backward.makeProbCache(s, f);
		double[][] forwardB = backward.makeForward(fSize, sSize, probCacheB);
		double[][] backwardB = backward.makeBackward(fSize, sSize, probCacheB);
		double likelihoodB = backward.makeLikelihood(forwardB[0], backwardB[0]);

		// Original caches
		double[][] probCacheFP = probCacheF.clone();
		double[][] probCacheBP = probCacheB.clone();

		double[][] posteriorsF = forward.makePosterior(forwardF, backwardF,
				likelihoodF);
		double[][] posteriorsB = backward.makePosterior(forwardB, backwardB,
				likelihoodB);

		// Initial Posteriors add them to alignment and print
		for (int fi = 0; fi < f.length; fi++) {
			for (int si = 0; si < s.length; si++) {
				fa.addPosteriro(si, fi, posteriorsF[fi][si]);
			}
		}

		for (int fi = 0; fi < f.length; fi++) {
			for (int si = 0; si < s.length; si++) {
				ba.addPosteriro(fi, si, posteriorsB[si][fi]);
			}
		}

		System.out.println("Sentence nr" + sentenceNumber);
		System.out.println("Before Projection" + sentenceNumber);
		AlignerOutputLatex.outputWithGoldPostLatex(fa, forwardCorpus.getGold()
				.get(sentenceNumber), forwardCorpus, System.out, "");
		System.out.println();
		System.out.println();
		AlignerOutputLatex.outputWithGoldPostLatex(ba, backwardCorpus.getGold()
				.get(sentenceNumber), backwardCorpus, System.out, "");
		System.out.println();
		System.out.println();
		double[][] phi = new double[fSize][sSize];
		for (int grI = 0; grI < _projectionIterations; grI++) {

			double lRate = 10.0 / (10.0 + grI);
			double totalViolation = 0;
			for (int fi = 0; fi < fSize; fi++) {
				for (int si = 0; si < sSize; si++) {
					// Update phi and lambda
					double ePsi = posteriorsF[fi][si] - posteriorsB[si][fi];
					double phiUpdate = lRate * (-ePsi);
					totalViolation += Math.abs(-ePsi);
					double newlambda = Math.abs(phi[fi][si]) - lRate * epsilon;
					double newPhi = phiUpdate + phi[fi][si];
					// Project
					if (-newPhi <= newlambda && newlambda <= newPhi) {
						newPhi = (newPhi + newlambda) / 2;
						newlambda = newPhi;
					} else if (newPhi <= newlambda && newlambda <= -newPhi) {
						newPhi = (newPhi - newlambda) / 2;
						newlambda = -newPhi;
					} else {
						newPhi = 0;
						newlambda = 0;
					}
					probCacheF[fi][si] = probCacheFP[fi][si] * Math.exp(newPhi);
					probCacheB[si][fi] = probCacheBP[si][fi]
							* Math.exp(-newPhi);
					phi[fi][si] = newPhi;
				}
			}
			// Renormalizing
			forwardF = forward.makeForward(sSize, fSize, probCacheF);
			backwardF = forward.makeBackward(sSize, fSize, probCacheF);
			likelihoodF = forward.makeLikelihood(forwardF[0], backwardF[0]);
			posteriorsF = forward.makePosterior(forwardF, backwardF,
					likelihoodF);

			forwardB = backward.makeForward(fSize, sSize, probCacheB);
			backwardB = backward.makeBackward(fSize, sSize, probCacheB);
			likelihoodB = backward.makeLikelihood(forwardB[0], backwardB[0]);
			posteriorsB = backward.makePosterior(forwardB, backwardB,
					likelihoodB);
		}

		for (int fi = 0; fi < f.length; fi++) {
			for (int si = 0; si < s.length; si++) {
				fa.addPosteriro(si, fi, posteriorsF[fi][si]);
			}
		}

		for (int fi = 0; fi < f.length; fi++) {
			for (int si = 0; si < s.length; si++) {
				ba.addPosteriro(fi, si, posteriorsB[si][fi]);
			}
		}
		AlignerOutputLatex.outputWithGoldPostLatex(fa, forwardCorpus.getGold()
				.get(sentenceNumber), forwardCorpus, System.out, "");
		System.out.println();
		System.out.println();
		AlignerOutputLatex.outputWithGoldPostLatex(ba, backwardCorpus.getGold()
				.get(sentenceNumber), backwardCorpus, System.out, "");
		System.out.println();
		System.out.println();
		System.out.println("Ended sentence " + sentenceNumber);
	}
	*/

	public static void main(String[] args) throws IOException {
		String corpusDescription = args[0];
		int size = Integer.parseInt(args[1]); // 100k
		int maxSentenceSize = Integer.parseInt(args[2]); // 40
		int numberIterations = Integer.parseInt(args[3]); // 5
		int projectionIterations = Integer.parseInt(args[4]); 
		// System.out.println("Corpus "+corpusName);
		System.out.println("Size " + size);
		System.out.println("Max Sentence size " + maxSentenceSize);
		System.out.println("Number of iterations " + numberIterations);
		System.out.println("Number of projection iterations " + projectionIterations);

		BilingualCorpus corpus = BilingualCorpus.getCorpusFromFileDescription(
				corpusDescription, size, maxSentenceSize);
		BilingualCorpus backwardCorpus = corpus.reverse();
		
		AgreementM1 m1 = new AgreementM1(corpus,backwardCorpus,0,projectionIterations);
		m1.train(numberIterations);

		AgreementHMMWithClone mhmm = new AgreementHMMWithClone(corpus, backwardCorpus,m1.forward._tb,m1.backward._tb,0,projectionIterations);
		mhmm.train(numberIterations);

		System.out.println("Done with training");

		AlignmentsSet sa = mhmm.viterbiAlignments(BilingualCorpus.TEST_CORPUS);
		AlignmentsSet gold = corpus.getGold();
		System.out.println("Gold size" + gold.size() + " viterbi size "
				+ sa.size());
		Evaluation eval2 = AlignmentEvaluator.evaluate(sa, corpus.getGold());
		System.out.println("Segment viterbi " + eval2);

		float tresh = mhmm.tuneTreshholdAER(BilingualCorpus.DEV_CORPUS,false);
		AlignmentsSet sa2 = mhmm.posteriorAlignments(
				BilingualCorpus.TEST_CORPUS, tresh,false,false);
		Evaluation eval22 = AlignmentEvaluator.evaluate(sa2, corpus.getGold());
		System.out.println("Posterior decoding " + eval22);
	}
	
}
